package com.gitee.luosl.cola.util;

import com.gitee.luosl.cola.lang.Tuple;
import com.gitee.luosl.cola.lang.Tuple2;
import com.gitee.luosl.cola.match.cases.Case;
import com.gitee.luosl.cola.stream.EntryStreamWrapper;
import com.gitee.luosl.cola.stream.StreamWrapper;

import java.util.*;
import java.util.function.BiFunction;
import java.util.function.Function;

import static com.gitee.luosl.cola.match.pattern.Pa.*;

public class HashJoiner {

    private static <T1, T2, K, R> R doHashJoin(Collection<T1> data1,
                                               Collection<T2> data2,
                                               Function<T1, K> keyExtractor1,
                                               Function<T2, K> keyExtractor2,
                                               BiFunction<EntryStreamWrapper<T1, K>, Map<K, List<T2>>, R> resultMapper) {
        Map<K, List<T2>> data2Map = StreamWrapper.wrap(data2).groupBy(keyExtractor2);
        EntryStreamWrapper<T1, K> data1Stream = StreamWrapper.wrapEntry(data1, Function.identity(), keyExtractor1);
        return resultMapper.apply(data1Stream, data2Map);
    }

    public static <T1, T2, K> List<Tuple2<T1, T2>> innerJoin(Collection<T1> data1,
                                                             Collection<T2> data2,
                                                             Function<T1, K> keyExtractor1,
                                                             Function<T2, K> keyExtractor2){
       return doHashJoin(
               data1, data2, keyExtractor1, keyExtractor2,
               (stream, map) -> stream.filter((ignore, k) -> map.containsKey(k))
                       .mapValues(map::get)
                       .flatMap((item, list) -> StreamWrapper.wrapEntry(list, ignore -> item, Function.identity()))
       ).toList();
    }

    public static <T, K> List<Tuple2<T, T>> innerJoin(Collection<T> data1,
                                                      Collection<T> data2,
                                                      Function<T, K> keyExtractor){
        return innerJoin(data1, data2, keyExtractor, keyExtractor);
    }

    public static <T1, T2, K> List<Tuple2<T1, Optional<T2>>> leftJoin(Collection<T1> data1,
                                                                      Collection<T2> data2,
                                                                      Function<T1, K> keyExtractor1,
                                                                      Function<T2, K> keyExtractor2){
        return doHashJoin(
                data1, data2, keyExtractor1, keyExtractor2,
                (stream, map) -> stream.mapValues(k -> Optional.ofNullable(map.get(k))).flatMap(
                        Case.of(pTuple(p(), pSome(pList())), (item, list) -> StreamWrapper.wrapEntry(list, ignore -> item, Optional::of)),
                        Case.of(pTuple(p(), pNone()), (item, ignore) -> StreamWrapper.ofEntry(Tuple.of(item, Optional.<T2>empty())))
                ).toList()
        );
    }

    public static <T, K> List<Tuple2<T, Optional<T>>> leftJoin(Collection<T> data1,
                                                               Collection<T> data2,
                                                               Function<T, K> keyExtractor){
        return leftJoin(data1, data2, keyExtractor, keyExtractor);
    }

    public static <T1, T2, K> List<Tuple2<Optional<T1>, T2>> rightJoin(Collection<T1> data1,
                                                                      Collection<T2> data2,
                                                                      Function<T1, K> keyExtractor1,
                                                                      Function<T2, K> keyExtractor2){
        Map<K, List<T1>> grouped = StreamWrapper.wrap(data1).groupBy(keyExtractor1);
        return StreamWrapper.wrap(data2)
                .mapToEntry(keyExtractor2, Function.identity())
                .mapKeys(k -> Optional.ofNullable(grouped.get(k)))
                .flatMap(
                        Case.of(pTuple(pSome(pList()), p()), (list, item) -> StreamWrapper.wrapEntry(list, Optional::of, ignore -> item)),
                        Case.of(pTuple(pNone(), p()), (ignore, item) -> StreamWrapper.ofEntry(Tuple.of(Optional.<T1>empty(), item)))
                ).toList();
    }

    public static <T, K> List<Tuple2<Optional<T>, T>> rightJoin(Collection<T> data1,
                                                                Collection<T> data2,
                                                                Function<T, K> keyExtractor){
        return rightJoin(data1, data2, keyExtractor, keyExtractor);
    }

    public static <T1, T2, K> List<Tuple2<Optional<T1>, Optional<T2>>> outerJoin(Collection<T1> data1,
                                                                                 Collection<T2> data2,
                                                                                 Function<T1, K> keyExtractor1,
                                                                                 Function<T2, K> keyExtractor2){
        Map<K, List<T1>> grouped1 = StreamWrapper.wrap(data1).groupBy(keyExtractor1);
        Map<K, List<T2>> grouped2 = StreamWrapper.wrap(data2).groupBy(keyExtractor2);

        EntryStreamWrapper<Optional<T1>, Optional<T2>> result1 = StreamWrapper.wrapEntry(data1, Function.identity(), keyExtractor1)
                .mapValues(k -> Optional.ofNullable(grouped2.get(k)))
                .flatMap(
                        Case.of(pTuple(p(), pSome(pList())), (T1 item, List<T2> list) -> StreamWrapper.wrapEntry(list, ignore -> Optional.of(item), Optional::of)),
                        Case.of(pTuple(p(), pNone()), (item, ignore) -> StreamWrapper.ofEntry(Tuple.of(Optional.of(item), Optional.empty())))
                );

        EntryStreamWrapper<Optional<T1>, Optional<T2>> result2 = StreamWrapper.wrapEntry(data2, keyExtractor2, Function.identity())
                .filterNot((k, v) -> grouped1.containsKey(k))
                .map(i -> Optional.empty(), Optional::of);

        return StreamWrapper.concat(result1, result2).toList();
    }

    public static <T, K> List<Tuple2<Optional<T>, Optional<T>>> outerJoin(Collection<T> data1,
                                                                          Collection<T> data2,
                                                                          Function<T, K> keyExtractor) {
        return outerJoin(data1, data2, keyExtractor, keyExtractor);
    }


}
